#include "Element.h"
#include "Quadrature.h"
#include <valarray>

using namespace std;

TriElement::TriElement()
{
	Element<2>::SetNDofs(n_dofs);
}

TriElement::TriElement(size_t _n_dofs)
{
	n_dofs = (_n_dofs >= 3 ? _n_dofs : 3);
	Element<2>::SetNDofs(n_dofs);
}

double TriElement::LambdaFunction(size_t i, const Point<2>& p) const
{
	double xi = p[0], eta = p[1], ret;
	switch (i) {
		case 0:
			ret = 1 - xi - eta;
			break;
		case 1:
			ret = xi;
			break;
		case 2:
			ret = eta;
			break;
		default:
			throw runtime_error("TriElement should have 0 <= i <= 2 !");
			break;
	}
	return ret;
}

valarray<double> TriElement::LambdaGradient(size_t i, const Point<2>& p) const
{
	valarray<double> ret(2);
	switch (i) {
		case 0:
			ret = valarray<double>({-1.0, -1.0});
			break;
		case 1:
			ret = valarray<double>({1.0, 0.0});
			break;
		case 2:
			ret = valarray<double>({0.0, 1.0});
			break;
		default:
			throw runtime_error("TriElement should have 0 <= i <= 2 !");
			break;
	}
	return ret;
}

Point<2> TriElement::LocalToGlobal(const Point<2>& p) const
{
	Point<2> global_point({0.0, 0.0});
	const valarray<Point<2> > &global_dofs = Element<2>::GetDofs();
	for (size_t i = 0; i < 3; ++i)
		global_point += global_dofs[i] * LambdaFunction(i, p);
	return global_point;
}

double TriElement::GlobalToLocalJacobiDet(const Point<2>& p) const
{
	const valarray<Point<2> > &dofs = Element<2>::GetDofs();
	return (dofs[1][0] - dofs[0][0]) * (dofs[2][1] - dofs[0][1])
		 - (dofs[1][1] - dofs[0][1]) * (dofs[2][0] - dofs[0][0]);
}

valarray<double> TriElement::LocalToGlobalJacobiTimes(const Point<2>& p, const valarray<double>& v) const
{
	const valarray<Point<2> > &dofs = Element<2>::GetDofs();
	double det = GlobalToLocalJacobiDet(p);
	valarray<double> ret(2);
	ret[0] = (dofs[2][1] - dofs[0][1]) * v[0] + (dofs[0][1] - dofs[1][1]) * v[1];
	ret[1] = (dofs[0][0] - dofs[2][0]) * v[0] + (dofs[1][0] - dofs[0][0]) * v[1];
	return ret / det;
}




double P1Element::BasisFunction(size_t i, const Point<2>& p) const
{
	double xi = p[0], eta = p[1], ret;
    switch (i) {
		case 0:
			ret = 1 - xi - eta;
			break;
		case 1:
			ret = xi;
			break;
		case 2:
			ret = eta;
			break;
		default:
			throw runtime_error("P1Element should have 0 <= i <= 2 !");
			break;

    }
	return ret;
}

valarray<double> P1Element::BasisGradient(size_t i, const Point<2>& p) const
{
	valarray<double> ret;
	switch (i) {
		case 0:
			ret = valarray<double>({-1.0, -1.0});
			break;
		case 1:
			ret = valarray<double>({1.0, 0.0});
			break;
		case 2:
			ret = valarray<double>({0.0, 1.0});
			break;
		default:
			throw runtime_error("P1Element should have 0 <= i <= 2 !");
			break;

    }
	return ret;
}

valarray<Point<2> > P1Element::DofsIndex(size_t ele_index, size_t n) const
{
	valarray<Point<2> > dofs_index(3);
	if (ele_index % 2 == 0) {
		dofs_index[0][0] = (ele_index / 2) % n; 
		dofs_index[1][0] = dofs_index[0][0] + 1;
		dofs_index[2][0] = dofs_index[0][0];

		dofs_index[0][1] = (ele_index / 2) / n;
		dofs_index[1][1] = dofs_index[0][1];
		dofs_index[2][1] = dofs_index[0][1] + 1;
	} else {
		dofs_index[0][0] = ((ele_index - 1) / 2) % n + 1;
		dofs_index[1][0] = dofs_index[0][0];
		dofs_index[2][0] = dofs_index[0][0] - 1;

		dofs_index[0][1] = ((ele_index - 1) / 2) / n;
		dofs_index[1][1] = dofs_index[0][1] + 1;
		dofs_index[2][1] = dofs_index[0][1] + 1;
	}
	return dofs_index;
}

size_t P1Element::SerialNumOfIndex(double i, double j, size_t n) const
{
	return static_cast<size_t>((n - 1) * (j - 1) + (i - 1)); 
}

size_t P1Element::GetNodeNum(size_t n) const
{
	return (n - 1) * (n - 1);
}


P2Element::P2Element() : TriElement(6) {}

double P2Element::BasisFunction(size_t i, const Point<2>& p) const
{
	double xi = p[0], eta = p[1], ret;
	switch (i) {
		case 0:
			ret = (1 - xi - eta) * (1.0 - 2.0 * xi - 2.0 * eta);
			break;		
		case 1:
			ret = xi * (2.0 * xi - 1.0);
			break;
		case 2:
			ret = eta * (2.0 * eta - 1.0);
			break;
		case 3:
			ret = 4.0 * xi * eta;
			break;
		case 4:
			ret = 4.0 * (1.0 - xi - eta) * eta;
			break;
		case 5:
			ret = 4.0 * (1.0 - xi - eta) * xi;
			break;
		default:
			throw runtime_error("P2Element should have 0 <= i <= 5 !");
			break;
    }
	return ret;
}

valarray<double> P2Element::BasisGradient(size_t i, const Point<2>& p) const
{
	double xi = p[0], eta = p[1];
    double phi_xi, phi_eta;

    switch (i) {
		case 0:
			phi_xi = phi_eta = -3.0 + 4.0 * xi + 4.0 * eta;
			break;		
		case 1:
			phi_xi = 4.0 * xi - 1.0;
			phi_eta = 0.0;
			break;
		case 2:
			phi_xi = 0.0;
			phi_eta = 4.0 * eta - 1.0;
			break;
		case 3:
			phi_xi = 4.0 * eta;
			phi_eta = 4.0 * xi;
			break;
		case 4:
			phi_xi = -4.0 * eta;
			phi_eta = 4.0 - 4.0 * xi - 8.0 * eta;
			break;
		case 5:
			phi_xi = 4.0 - 4.0 * eta - 8.0 * xi;
			phi_eta = -4.0 * xi;
			break;
		default:
			throw runtime_error("P2Element should have 0 <= i <= 5 !");
			break;
    }
	return valarray<double>({phi_xi, phi_eta});
}

valarray<Point<2> > P2Element::DofsIndex(size_t ele_index, size_t n) const
{
	valarray<Point<2> > dofs_index(6);
	if (ele_index % 2 == 0) {
		dofs_index[0][0] = (ele_index / 2) % n; 
		dofs_index[1][0] = dofs_index[0][0] + 1.0;
		dofs_index[2][0] = dofs_index[0][0];
		dofs_index[3][0] = dofs_index[0][0] + 0.5;
		dofs_index[4][0] = dofs_index[0][0];
		dofs_index[5][0] = dofs_index[0][0] + 0.5;

		dofs_index[0][1] = (ele_index / 2) / n;
		dofs_index[1][1] = dofs_index[0][1];
		dofs_index[2][1] = dofs_index[0][1] + 1.0;
		dofs_index[3][1] = dofs_index[0][1] + 0.5;
		dofs_index[4][1] = dofs_index[0][1] + 0.5;
		dofs_index[5][1] = dofs_index[0][1];
	} else {
		dofs_index[0][0] = ((ele_index - 1) / 2) % n + 1;
		dofs_index[1][0] = dofs_index[0][0];
		dofs_index[2][0] = dofs_index[0][0] - 1.0;
		dofs_index[3][0] = dofs_index[0][0] - 0.5;
		dofs_index[4][0] = dofs_index[0][0] - 0.5;
		dofs_index[5][0] = dofs_index[0][0];

		dofs_index[0][1] = ((ele_index - 1) / 2) / n;
		dofs_index[1][1] = dofs_index[0][1] + 1.0;
		dofs_index[2][1] = dofs_index[0][1] + 1.0;
		dofs_index[3][1] = dofs_index[0][1] + 1.0;
		dofs_index[4][1] = dofs_index[0][1] + 0.5;
		dofs_index[5][1] = dofs_index[0][1] + 0.5;
	}
	return dofs_index;
}

size_t P2Element::SerialNumOfIndex(double i, double j, size_t n) const
{
	return static_cast<size_t>((2 * n - 1) * (2 * j - 1) + (2 * i - 1));
}

size_t P2Element::GetNodeNum(size_t n) const
{
	return (2 * n - 1) * (2 * n - 1);
}


